# -*- encoding: utf-8 -*-
'''
@File    :   gmm.py
@Time    :   2021/11/25 10:21
@Author  :   ZhangChaoYang
@Desc    :   Gaussian Mixture Model
'''
import os

import numpy as np
import tensorflow as tf
import traceback


class GMM(object):
    def __init__(self):
        '''
        K: component（簇）的个数
        d: 特征个数
        N: 样本个数
        phi: 各个component的概率，shape=(K,)
        mu: 各个component的均值，shape=(K,d)
        sigma: 各个component内各特征之间的协方差，shape=(K,d,d)
        '''
        self.phi = self.mu = self.sigma = None
        self.threshold = 0

    def fit(self, z, gamma):
        '''
        z: 样本，shape=(N,d)
        gamma: 样本属于各个component的概率，shape=(N,K)
        '''
        self.phi = tf.reduce_mean(gamma, axis=0)
        gamma_sum = tf.reduce_sum(gamma, axis=0)
        self.mu = tf.einsum('ik,il->kl', gamma, z) / gamma_sum[:, None]
        z_centered = tf.sqrt(gamma[:, :, None]) * (z[:, None, :] - self.mu[None, :, :])
        self.sigma = tf.einsum('ikl,ikm->klm', z_centered, z_centered) / gamma_sum[:, None, None]
        # 提前计算协方差矩阵的cholesky分解
        n_features = z.shape[1]
        min_vals = tf.linalg.diag(tf.ones(n_features, dtype=tf.float32)) * 1e-4
        self.L = tf.linalg.cholesky(self.sigma + min_vals[None, :, :])
        # self.L = tf.linalg.cholesky(self.sigma)

    def energy(self, z):
        '''
        计算z中每一个样本的sample energy(负的似然)
        '''
        z_centered = z[:, None, :] - self.mu[None, :, :]  # (N,K,d)
        try:
            v = tf.linalg.triangular_solve(self.L, tf.transpose(z_centered, [1, 2, 0]))  # (K,d,N)
        except:
            # 矩阵不可逆，这里需要做重大事故报警
            traceback.print_exc(-1)
            raise Exception("模型训练失败")
        else:
            # log(det(Sigma)) = 2 * sum[log(diag(L))]
            log_det_sigma = 2.0 * tf.reduce_sum(tf.math.log(tf.linalg.diag_part(self.L)), axis=1)

            # To calculate energies, use "log-sum-exp" (different from orginal paper)
            d = z.get_shape().as_list()[1]
            logits = tf.math.log(self.phi[:, None]) - 0.5 * (tf.reduce_sum(tf.square(v), axis=1)
                                                             + d * tf.math.log(2.0 * np.pi) + log_det_sigma[:, None])
            energies = - tf.reduce_logsumexp(logits, axis=0)
            return energies

    def cov_diag_loss(self):
        diag_loss = tf.reduce_mean(tf.divide(1, tf.linalg.diag_part(self.sigma)))
        return diag_loss

    def save(self, model_file):
        if not os.path.exists(model_file):
            os.makedirs(model_file)
        np.save(os.path.join(model_file, "phi"), self.phi.numpy())
        np.save(os.path.join(model_file, "mu"), self.mu.numpy())
        np.save(os.path.join(model_file, "L"), self.L.numpy())
        np.save(os.path.join(model_file, "threshold"), np.asarray([self.threshold]))

    def load(self, model_file):
        self.phi = tf.convert_to_tensor(np.load(os.path.join(model_file, "phi.npy")))
        self.mu = tf.convert_to_tensor(np.load(os.path.join(model_file, "mu.npy")))
        self.L = tf.convert_to_tensor(np.load(os.path.join(model_file, "L.npy")))
        self.threshold = np.load(os.path.join(model_file, "threshold.npy"))[0]
